Elastic Diffusion for Wan model by huvunvidia · Pull Request #34 · NVlabs/FastGen

huvunvidia · 2026-07-16T03:50:45Z

No description provided.

Design doc for porting the Elastification (Flextron) one-shot NAS compression algorithm from Megatron-LM to FastGen, targeting Wan. Covers scope decisions, file structure, per-file design, test plan, and open items ahead of implementation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Removes algorithm summary, open decisions, test plan, implementation order, non-goals, and appendices — leaving locked-in decisions, Wan integration points, file structure, per-file design notes, and inference. Also switches distillation to external frozen teacher and scopes v1 to Wan 2.1 T2V (drops I2V / TI2V mentions). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Ports Megatron-LM's FlextronConfig and FlextronRouter to FastGen: - config.py: FlextronConfig attrs dataclass, section/field order mirrors flextron_config.py. Wan-scope drops (Mamba, MoE, memory, basemodel_type); Wan additions (lr_mult_router, train_iters, num_layers, num_mlp_blocks). - router.py: FlextronRouter class, method order + naming mirror hybrid_flex_router.py. Includes DP-seeded Gumbel-softmax with RNG isolation, per-axis forwards (mlp/emb/skipping), tau decay, optional scaler schedule, and interpolated one-hot for intermediate budgets at inference. Deviations from the reference are noted inline. One Megatron bug (unbound `scale` in mlp_forward when scaler is disabled) is fixed here rather than mirrored; PLAN.md flags this as a TODO to verify whether Megatron's actual training configs enable the scaler. Runtime-verified against 13 checks: construction, gate shapes + init std, forward-shape, RNG isolation, DP-seed determinism, tau decay, budget interpolation, gradient flow through non-invariant loss, add_skipping mask, flex_hetero_ffn per-block choice list. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

48 test cases across 8 groups (TestConstruction, TestForward, TestDPGumbelSoftmax, TestTauDecay, TestGradientFlow, TestAddSkipping, TestFlexHeteroFfn, TestScalerOffPath). Logic mirrors 13 ad-hoc runtime checks already passed against router.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Ports Megatron-LM's flex_budget_utils::get_num_parameters, adapted for Wan diffusion transformers. Accepts int or torch.Tensor widths for hidden_size / mlp_hidden_size so the same formula backs both hard sub-net sizing and the differentiable budget loss on soft router outputs. Wan-scope drops: no Mamba, no MoE, no hybrid_pattern, no kv_channels/GQA split, no memory-mode. Wan-scope adds: text_encoder_dim (attn2 K/V input), patch_dim + in/out_channels (patch-embed + proj-out), time_freq_dim (sinusoidal-embedding intermediate). Per-block formula: attn1 (4*h^2) + attn2 (2*h^2 + 2*text_dim*h) + ffn (2*h*mlp) + adaLN (6*h) + norms. Optional per-layer skip probability tensor weights each block's contribution by (1 - skip_prob). Test coverage (16 pytest cases): return type, positive value, halving scaling behavior on both axes, layer count linearity, all-zeros/all-ones skip edge cases, wrong-shape validation, and end-to-end gradient flow through soft router probs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Structure now closely mirrors Megatron's flex_budget_utils.py: - Same argument names (hidden_size, ffn_hidden_size, num_attention_heads) - norm_multiplier constant - Hetero detection via isinstance + .shape[0] != 1 - FFN branch mirrors moe_all structure (collapsed — no MoE) - ATT branch computes self-attn without GQA split (Wan doesn't use GQA) - Wan-specific sections (cross-attn attn2, adaLN, patch/proj boilerplate) tagged inline - Per-block loop mirrors hybrid_pattern iteration structure Delete tests/elastification/test_budget_utils.py: the tests only checked relative behavior (halving hidden reduces params, gradients flow through PyTorch ops) with no verification of absolute correctness. The right test pattern (per Megatron's test_flex_budget_utils.py) is exact-integer equality vs a hand-computed reference; that will land in a follow-up along with a Wan-model-based end-to-end check. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

greptile-apps · 2026-07-16T03:56:01Z

Greptile Summary

This PR introduces the first three files of the Elastification / Flextron elastic diffusion method for the Wan model, porting Megatron-LM's FlextronRouter and analytical parameter-count utility into FastGen, along with FlextronConfig and a comprehensive router test suite.

router.py implements the Gumbel-softmax router with DP-seeded noise, budget interpolation, and a scaler schedule, and deliberately fixes a Megatron NameError in mlp_forward; however, emb_forward silently drops the normalize_router_logits normalization that mlp_forward applies, the scaler tensor is indexed without bounds checking (crashing if curr_iteration ≥ train_iters), and the constructor crashes with TypeError when budget_list=None (the config default).
budget_utils.py provides the Wan-specific analytical param count for differentiable budget loss; the ScalarOrTensor type alias claims float support but the het-FFN detection branch calls .shape[0] on a plain Python float, raising AttributeError.
The test suite is well-structured with good axis coverage but does not exercise the normalize_router_logits path on emb_forward or any None-default config scenario.

Confidence Score: 3/5

The router has three independent crash/correctness paths that will be hit in real training runs.

The emb_forward method silently misroutes the embedding axis whenever normalize_router_logits=True, producing inconsistent routing invisible in logs. The scaler index goes out of bounds on the final training iteration or during eval with a large iteration count. The budget_list=None default lets a misconfigured config reach deep into the router constructor before crashing. All three affect the core training loop.

fastgen/methods/elastification/router.py needs the most attention — all three constructor/forward bugs live there. fastgen/methods/elastification/budget_utils.py needs a one-line type-guard fix.

Important Files Changed

Filename	Overview
fastgen/methods/elastification/router.py	Core FlextronRouter implementation with three bugs: `emb_forward` ignores `normalize_router_logits`, scaler tensor lacks bounds checking, and `budget_list=None` crashes the constructor. Also contains a dead `repeat_interleave(repeats=1)` no-op in `skipping_forward`.
fastgen/methods/elastification/budget_utils.py	Analytical Wan parameter-count function; `ScalarOrTensor` type alias claims `float` support but the het-FFN branch crashes with `AttributeError` when a Python `float` is passed for `ffn_hidden_size`.
fastgen/methods/elastification/config.py	FlextronConfig attrs dataclass with validation; `budget_list` defaults to `None` rather than a required field, enabling a crash path in the router constructor that could be caught here instead.
tests/elastification/test_router.py	Thorough router unit tests covering construction, forward paths, DP-seeded Gumbel, tau decay, gradients, skip, hetero-FFN, and the Megatron NameError fix; all fixtures provide valid configs so the `budget_list=None` and `emb_forward`/`normalize_router_logits` gaps are not exercised.
fastgen/methods/elastification/PLAN.md	Design document describing the Elastification port from Megatron-LM; no executable code, well-written with explicit callouts for known Megatron quirks and scope reductions.
tests/elastification/init.py	Empty test package init file; no issues.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Router as FlextronRouter.forward()
    participant BudgetMap as budget_map lookup
    participant MLP as mlp_forward()
    participant EMB as emb_forward()
    participant SKIP as skipping_forward()
    participant Gumbel as _dp_gumbel_softmax()

    Caller->>Router: forward(budget, curr_iteration)
    Router->>Router: "hard_sample = random() > hard_sample_th"
    Router->>Router: "tau = tau_init * tau_decay^curr_iter"

    alt budget in budget_map
        Router->>BudgetMap: one_hot lookup
    else "budget == 1.0 fallback"
        Router->>BudgetMap: one_hot for max(budget_map)
    else interpolated budget
        Router->>Router: bucketize + linear blend of two one-hots
    end

    Router->>MLP: mlp_forward(curr_iter, budget_tensor, ...)
    MLP->>MLP: gate_mlp forward pass
    MLP->>MLP: apply scaler + normalize_router_logits branch
    MLP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    note over Gumbel: save/restore CPU+CUDA RNG state
    Gumbel-->>MLP: soft_probs
    MLP-->>Router: (soft_probs, chosen_int)

    opt "add_skipping=True"
        Router->>SKIP: skipping_forward(curr_iter, budget_tensor, ...)
        SKIP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
        Gumbel-->>SKIP: soft_probs
        SKIP->>SKIP: build per-layer 0/1 skip mask
        SKIP-->>Router: (soft_probs, skip_mask)
    end

    Router->>EMB: emb_forward(curr_iter, budget_tensor, ...)
    EMB->>EMB: gate_emb forward pass
    note over EMB: normalize_router_logits ignored here
    EMB->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    Gumbel-->>EMB: soft_probs
    EMB-->>Router: (soft_probs, chosen_int)

    Router->>Router: "fwd_pass_count += 1"
    Router-->>Caller: (mlp_out, skip_out or None, emb_out)

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant Router as FlextronRouter.forward()
    participant BudgetMap as budget_map lookup
    participant MLP as mlp_forward()
    participant EMB as emb_forward()
    participant SKIP as skipping_forward()
    participant Gumbel as _dp_gumbel_softmax()

    Caller->>Router: forward(budget, curr_iteration)
    Router->>Router: "hard_sample = random() > hard_sample_th"
    Router->>Router: "tau = tau_init * tau_decay^curr_iter"

    alt budget in budget_map
        Router->>BudgetMap: one_hot lookup
    else "budget == 1.0 fallback"
        Router->>BudgetMap: one_hot for max(budget_map)
    else interpolated budget
        Router->>Router: bucketize + linear blend of two one-hots
    end

    Router->>MLP: mlp_forward(curr_iter, budget_tensor, ...)
    MLP->>MLP: gate_mlp forward pass
    MLP->>MLP: apply scaler + normalize_router_logits branch
    MLP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    note over Gumbel: save/restore CPU+CUDA RNG state
    Gumbel-->>MLP: soft_probs
    MLP-->>Router: (soft_probs, chosen_int)

    opt "add_skipping=True"
        Router->>SKIP: skipping_forward(curr_iter, budget_tensor, ...)
        SKIP->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
        Gumbel-->>SKIP: soft_probs
        SKIP->>SKIP: build per-layer 0/1 skip mask
        SKIP-->>Router: (soft_probs, skip_mask)
    end

    Router->>EMB: emb_forward(curr_iter, budget_tensor, ...)
    EMB->>EMB: gate_emb forward pass
    note over EMB: normalize_router_logits ignored here
    EMB->>Gumbel: _dp_gumbel_softmax(logits, tau, hard)
    Gumbel-->>EMB: soft_probs
    EMB-->>Router: (soft_probs, chosen_int)

    Router->>Router: "fwd_pass_count += 1"
    Router-->>Caller: (mlp_out, skip_out or None, emb_out)

_{Reviews (1): Last reviewed commit: "Rewrite budget_utils.py to mirror Megatr..." | Re-trigger Greptile}

greptile-apps · 2026-07-16T03:56:05Z

+    def emb_forward(self, curr_iteration, budget_tensor, device, dtype, tau, hard_sample):
+
+        # `[0]` unpacks removed vs Megatron — see comment in mlp_forward.
+        router_emb_logits1 = self.gate_emb[0](budget_tensor)
+        router_emb_logits2 = self.gate_emb[1](router_emb_logits1)
+        router_emb_logits = self.gate_emb[2](router_emb_logits2).flatten()
+        if self.scaler is not None:
+            scale = self.scaler[curr_iteration].to(device=device, dtype=dtype)
+            router_emb_logits = scale * router_emb_logits
+
+        router_emb_logits = self._dp_gumbel_softmax(
+            router_emb_logits, tau=tau, hard=hard_sample, curr_iteration=curr_iteration
+        )
+        _, choices_emb = torch.topk(router_emb_logits, 1, dim=-1)
+
+        return (router_emb_logits, self.config.emb_int_list[choices_emb.item()])


emb_forward silently ignores normalize_router_logits

mlp_forward fully implements the normalize_router_logits config flag — normalizing logits by their std before Gumbel-softmax when enabled. emb_forward has no such branch: it only applies the scaler when self.scaler is not None and passes raw logits otherwise, completely ignoring normalize_router_logits. When normalize_router_logits=True, the MLP axis logits are std-normalized while the EMB axis logits are not, producing an asymmetric routing that diverges from the intended configuration and from the Megatron reference.

greptile-apps · 2026-07-16T03:56:06Z

+        if self.scaler is not None:
+            scale = self.scaler[curr_iteration].to(device=device, dtype=dtype)


Scaler tensor indexed without bounds checking

self.scaler is built by torch.linspace(..., steps=train_iters), giving a tensor of length train_iters. All three axis forward methods (mlp_forward, emb_forward, skipping_forward) index it as self.scaler[curr_iteration] without clamping. If curr_iteration >= train_iters — for example, at the final iteration of training (off-by-one), during fine-tuning from a checkpoint with a longer schedule, or in eval mode with a large iteration counter — this raises an IndexError. The fix is to clamp curr_iteration to [0, len(self.scaler) - 1] before indexing.

greptile-apps · 2026-07-16T03:56:07Z

+        self.input_dim = len(self.config.budget_list)
+        self.n_dim = self.config.router_inter_dim
+        self.budget_map = {
+            item: torch.tensor(idx)
+            for idx, item in enumerate(self.config.budget_list)
+        }


budget_list=None causes TypeError at construction

FlextronConfig.budget_list defaults to None (line 56 of config.py). FlextronRouter.__init__ calls len(self.config.budget_list) and iterates over it unconditionally on lines 60–65. Constructing FlextronRouter with a default FlextronConfig() (or any config that omits budget_list) crashes immediately with TypeError: object of type 'NoneType' has no len(). A guard — either in FlextronConfig.__attrs_post_init__ that asserts budget_list is not None, or in FlextronRouter.__init__ — would surface this at config-construction time with a clear message rather than a cryptic stack trace mid-constructor.

greptile-apps · 2026-07-16T03:56:08Z

+        router_skip_layer_logits = torch.repeat_interleave(
+            router_skip_layer_logits, repeats=1, dim=0
+        )


torch.repeat_interleave(x, repeats=1, dim=0) returns the input tensor unchanged — repeating each element exactly once is a no-op. This dead code (inherited from Megatron, where it likely served a different purpose) adds a needless allocation on every skipping_forward call.

Suggested change

router_skip_layer_logits = torch.repeat_interleave(

router_skip_layer_logits, repeats=1, dim=0

)

# repeat_interleave(repeats=1) is a no-op; removed.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

greptile-apps · 2026-07-16T03:56:09Z

+    if isinstance(ffn_hidden_size, int):
+        flex_hetero_ffn = False
+    else:
+        flex_hetero_ffn = ffn_hidden_size.shape[0] != 1


float type triggers AttributeError in the het-FFN detection branch

The type alias ScalarOrTensor = Union[int, float, torch.Tensor] advertises that ffn_hidden_size can be a Python float, but the branch at line 69 only handles int and torch.Tensor. If a caller passes a plain float (e.g., ffn_hidden_size=5000.0), isinstance(ffn_hidden_size, int) is False, so execution falls to ffn_hidden_size.shape[0], which raises AttributeError: 'float' object has no attribute 'shape'. Either update the type hint to Union[int, torch.Tensor] to reflect the actual contract, or add isinstance(ffn_hidden_size, (int, float)) in the first branch to handle both scalar cases.

Huy Vu and others added 6 commits July 13, 2026 13:35

greptile-apps Bot reviewed Jul 16, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Elastic Diffusion for Wan model#34

Elastic Diffusion for Wan model#34
huvunvidia wants to merge 6 commits into
NVlabs:mainfrom
huvunvidia:huvu/elastic_diffusion

huvunvidia commented Jul 16, 2026

Uh oh!

greptile-apps Bot commented Jul 16, 2026

Uh oh!

greptile-apps Bot Jul 16, 2026

Uh oh!

greptile-apps Bot Jul 16, 2026

Uh oh!

greptile-apps Bot Jul 16, 2026

Uh oh!

greptile-apps Bot Jul 16, 2026

Uh oh!

greptile-apps Bot Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		if self.scaler is not None:
		scale = self.scaler[curr_iteration].to(device=device, dtype=dtype)

Uh oh!

Conversation

huvunvidia commented Jul 16, 2026

Uh oh!

greptile-apps Bot commented Jul 16, 2026

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant